import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from scipy.integrate import quad

p = 0.55  # 间距长度,单位m
v_head = 1.0  # 头部速度

det_t=0.1       #时间间隙
times=int(1/det_t)
t_total = 1000  # 时间总量300s
sequen=t_total*times
cal_finaltime=t_total

num_sections = 223  # 节数
length_head = 3.41  # 头部长度
length_body = 2.20  # 身体部分长度
distance=0.275

sections_lengths = [length_head] + [length_body] * (num_sections - 1)
positions = np.zeros((sequen+times, num_sections+1,2))
velocities = np.zeros((sequen + times, num_sections))
positions_x_y=np.zeros((sequen + times, num_sections+1,2))

#R=A+B*THETA
a=p*16
b=p/2/np.pi

def calculate_R(theta):
    return a+b*theta
def x_y_cal(t,section_index):
    bu_jin=int(t*times)
    x,y=positions[bu_jin,section_index,0]*np.cos(positions[bu_jin,section_index,1]),positions[bu_jin,section_index,0]*np.sin(positions[bu_jin,section_index,1])
    return x,y


def calculate_position(t, section_index):
    index=int(t*times)
    r=0
    theta=0
    if section_index == 0:#公式推到
        theta=(np.sqrt(a**2/b+2*t)+a/np.sqrt(b))/np.sqrt(b)*(-1)+(2*a/b)
        r = calculate_R(theta)
    else:
        prev_r, prev_t = positions[index, section_index - 1,0],positions[index, section_index - 1,1]#由前一节点，板凳长度推导
        d=sections_lengths[section_index-1]-2*distance
        theta=prev_t+np.sqrt(d**2/(1+(prev_r**2))*b**2)
        r=calculate_R(theta)
        if r>a:
             r=theta=0
    if(r<=0):
        r=theta=0
    return r, theta

def calculate_velocity(t, section_index):
    index=int(t*times)
    if index <3:
        return 0
    # now_r = positions[index, section_index,0]
    # prev_t = positions[index-1, section_index,1]
    # now_t=positions[index, section_index,1]
    # v_ = (now_t-prev_t)/det_t*np.sqrt(b**2+now_r**2)

    prev_x,prev_y=positions_x_y[index-1,section_index,0],positions_x_y[index-1,section_index,1]
    if(prev_y==0 and prev_x==0):
        return 0
    now_x,now_y=positions_x_y[index,section_index,0],positions_x_y[index,section_index,1]
    v_=np.sqrt((now_x-prev_x)**2+(now_y-prev_y)**2)/det_t
    # v_, error = quad(calculate_R, prev_x, now_x)
    # v_=v_/2/det_t
    return v_

def plot_positions():
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.set_aspect('equal')
    for t in [0*times, 60*times, 120*times, 180*times, 240*times, 300*times]:
        ax.plot(positions_x_y[t, :,0], positions_x_y[t, :,1], label=f't={t}s')
    ax.scatter(positions_x_y[:, 0,0], positions_x_y[:, 0,1], color='red', s=50, label='trace of  loong head')
    ax.set_title('the trace of the team')
    ax.set_xlabel('x_position(m)')
    ax.set_ylabel('y_position (m)')
    ax.legend()
    plt.grid(True)
    plt.show()

def write_excel():
    result = pd.DataFrame(columns=["time", "section", "x_position", "y_position", "r_","theta","velocity"])
    data = []

    for k in np.arange(0, cal_finaltime, 1):
        for j in range(num_sections):
            i = int(k * times)
            data.append({
                "time": i,
                'section': j + 1,
                'x_position': positions_x_y[i, j, 0],
                'y_position': positions_x_y[i, j, 1],
                "r_":positions[i,j,0],
                "theta" :positions[i,j,1],
                'velocity': velocities[i, j]
            })
    print(1)
    result = pd.DataFrame(data)
    result.to_excel('result.xlsx', index=False)

for t in np.arange(0,cal_finaltime,det_t):
    for i in range(num_sections):
        index_arr=int(t*times)

        r_t_temp=calculate_position(t, i)
        positions[index_arr, i, 0], positions[index_arr, i, 1] = r_t_temp[0], r_t_temp[1]

        x_y_temp=x_y_cal(t,i)
        positions_x_y[index_arr,i,0],positions_x_y[index_arr,i,1]=x_y_temp[0],x_y_temp[1]

        velocities[index_arr, i] = calculate_velocity(t, i)
plot_positions()
#write_excel()
